package agent;

import javassist.*;
import transformers.AgentTransformer;
import transformers.LoaderTransformer;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.net.URL;
import java.net.URLDecoder;
import java.security.ProtectionDomain;
import java.util.jar.JarFile;

/**
 * Author: 4ut15m.4ny0neSec
 * Desc:
 * Agent的demo
 * agent要打包成jar而后加载运行，不能直接跑。
 * <p>
 * 打包方法：
 * idea -> Terminal -> cd到agentdemo的目录下 -> mvn clean install
 * 最后jar在agentdemo/target/目录下
 */
public class Agent {
    private static String agentPath;

    /**
     * 以 -javaagent方式加载agent时会执行premain方法
     * premain会在主程序的main方法执行之前执行
     * 一些方法我会注释，你想运行看看效果的时候自行取消注释
     */
    public static void premain(String args, Instrumentation instrumentation) {
//        System.out.println("This function is premain.");
//        printStack(instrumentation);
        /**
         * 下面两行是用javassist修改字节码的样例
         * */
//        ClassFileTransformer classFileTransformer = initJavasist();
//        instrumentation.addTransformer(classFileTransformer);
        /**
         * 下面两行是用asm修改字节码的样例
         * */
//        ClassFileTransformer classFileTransformer = initASM();
//        instrumentation.addTransformer(classFileTransformer);
        /**
         * 下面是hook ProcessBuilder的一个样例
         * 为防止程序访问不到我们添加的hook类，需要使用instrumentation.appendToBootstrapClassLoaderSearch()
         * 将我们的agent添加进BootstrapClassLoader中，否则后续可能会出现找不到类的问题。
         * */
//        try {
//            setLocalAgentPath();
//            instrumentation.appendToBootstrapClassLoaderSearch(new JarFile(agentPath));
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
//        ClassFileTransformer classFileTransformer = hookProcessBuilder();
//        instrumentation.addTransformer(classFileTransformer);
    }

    /**
     * 在以attach方式加载agent的时候会运行agentmain方法
     */
    public static void agentmain(String args, Instrumentation instrumentation) {
        System.out.println("This function is agentmain");
        printStack(instrumentation);
    }

    /**
     * Desc:
     * Javasist修改字节码，修改controller.Load类的字节码
     * 在sayHello方法的最后添加一行sout输出。
     */
    public static ClassFileTransformer initJavasist() {
        //Javasist修改字节码
        return new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                className = className.replaceAll("/", ".");
                if (className.equals("com.any0ne.loadDemo.Load")) {
                    ClassPool classPool = ClassPool.getDefault();
                    CtClass ctClass = null;
                    CtMethod ctMethod = null;
                    try {
                        ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
                        ctMethod = ctClass.getDeclaredMethod("sayHello");
                        ctMethod.insertBefore("{ System.out.println(\"4ut15m\"); }");
                        ctMethod.insertAfter("{ System.out.println(\"world\"); }");
                        return ctClass.toBytecode();
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
                return null;
            }
        };
    }

    public static ClassFileTransformer initASM() {
        return new LoaderTransformer();
    }

    public static ClassFileTransformer hookProcessBuilder() {
        return new AgentTransformer();
    }

    /**
     * 打印程序加载的所有类
     * 重点是怎么获取加载的类，不是打印
     */
    public static void printStack(Instrumentation instrumentation) {
        Class[] c = instrumentation.getAllLoadedClasses();
        for (Class aClass : c) {
            System.out.println(aClass);
        }
    }

    /**
     * 设置Agent的路径。
     */
    public static void setLocalAgentPath() {
        URL u = Agent.class.getProtectionDomain().getCodeSource().getLocation();
        try {
            agentPath = URLDecoder.decode(u.getFile().replace("+", "%2B"), "UTF-8");
        } catch (UnsupportedEncodingException e) {
            e.printStackTrace();
        }

        if (agentPath.lastIndexOf("/") != -1) {
            agentPath.substring(0, agentPath.lastIndexOf("/"));
        }
    }


}
